#!/usr/bin/env python3
"""
拓扑地图可视化工具
支持加载包含坐标和连接关系的拓扑节点，并进行可视化展示
"""

import yaml
import networkx as nx
import matplotlib.pyplot as plt
import sys
import os

def load_topology_yaml(file_path):
    """加载 YAML 拓扑文件"""
    if not os.path.exists(file_path):
        print(f"❌ 错误：文件 '{file_path}' 不存在！")
        sys.exit(1)
    
    with open(file_path, 'r', encoding='utf-8') as f:
        try:
            data = yaml.safe_load(f)
        except yaml.YAMLError as e:
            print(f"❌ YAML 解析错误: {e}")
            sys.exit(1)
    
    if not isinstance(data, dict) or 'node_list' not in data:
        print("❌ 错误：YAML 文件应包含 'node_list' 键")
        sys.exit(1)
    
    return data['node_list']

def parse_topology_nodes(nodes):
    """
    解析拓扑节点
    param 格式: [x, y, z, yaw, connected_nodes, node_type]
    - param[0]: x 坐标
    - param[1]: y 坐标
    - param[2]: z 坐标
    - param[3]: yaw 角度
    - param[4]: 连接的节点（逗号分隔的字符串）
    - param[5]: 节点类型
    """
    G = nx.Graph()
    positions = {}
    node_info = {}
    
    for node in nodes:
        if not isinstance(node, dict):
            print(f"⚠️  警告：跳过非字典项: {node}")
            continue
        
        node_name = node.get('node_name')
        if not node_name:
            print(f"⚠️  警告：节点缺少 'node_name' 字段，跳过")
            continue
        
        params = node.get('param', [])
        if not isinstance(params, list) or len(params) < 6:
            print(f"⚠️  警告：节点 '{node_name}' 的 param 格式不正确，跳过")
            continue
        
        # 提取节点信息
        x = float(params[0])
        y = float(params[1])
        z = float(params[2])
        yaw = float(params[3])
        connected_str = params[4]
        node_type = params[5]
        
        # 获取floor信息
        floor = node.get('floor', 1)
        
        # 添加节点
        G.add_node(node_name)
        positions[node_name] = (x, y)
        node_info[node_name] = {
            'position': (x, y, z),
            'yaw': yaw,
            'floor': floor,
            'type': node_type
        }
        
        # 解析连接的节点（处理逗号分隔的字符串）
        if connected_str and connected_str != '':
            # 去除可能的引号并分割
            connected_str = str(connected_str).strip("'\"")
            connected_nodes = [n.strip() for n in connected_str.split(',')]
            
            for neighbor in connected_nodes:
                if neighbor and neighbor != '':
                    G.add_edge(node_name, neighbor)
    
    return G, positions, node_info

def print_topology_info(G, positions, node_info):
    """打印拓扑信息"""
    print("\n" + "="*60)
    print("📊 拓扑网络信息")
    print("="*60)
    print(f"节点总数: {G.number_of_nodes()}")
    print(f"边总数: {G.number_of_edges()}")
    print("\n" + "-"*60)
    print("节点详细信息:")
    print("-"*60)
    
    for node in sorted(G.nodes()):
        info = node_info.get(node, {})
        pos = positions.get(node, (0, 0))
        neighbors = list(G.neighbors(node))
        
        print(f"\n🔹 节点: {node}")
        print(f"   位置: ({pos[0]:.2f}, {pos[1]:.2f})")
        if info:
            print(f"   类型: {info.get('type', 'N/A')}")
            print(f"   楼层: {info.get('floor', 'N/A')}")
            print(f"   偏航角: {info.get('yaw', 0):.2f} rad")
        print(f"   连接到: {', '.join(neighbors) if neighbors else '无'}")
        print(f"   连接数: {len(neighbors)}")

def visualize_topology(G, positions, node_info, title="拓扑地图网络"):
    """可视化拓扑网络"""
    plt.figure(figsize=(14, 10))
    
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Arial Unicode MS', 'SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 绘制边
    nx.draw_networkx_edges(G, positions, width=2.5, alpha=0.5, 
                          edge_color='gray', style='solid')
    
    # 根据节点类型设置不同颜色
    node_colors = []
    node_sizes = []
    for node in G.nodes():
        info = node_info.get(node, {})
        node_type = info.get('type', '1')
        
        # 根据类型分配颜色
        if node_type == '1':
            node_colors.append('lightblue')
            node_sizes.append(2500)
        elif node_type == '2':
            node_colors.append('lightgreen')
            node_sizes.append(2500)
        elif node_type == '3':
            node_colors.append('lightyellow')
            node_sizes.append(2500)
        else:
            node_colors.append('lightcoral')
            node_sizes.append(2500)
    
    # 绘制节点
    nx.draw_networkx_nodes(G, positions, 
                          node_color=node_colors,
                          node_size=node_sizes,
                          edgecolors='black',
                          linewidths=2.5,
                          alpha=0.9)
    
    # 绘制节点标签
    nx.draw_networkx_labels(G, positions, 
                           font_size=12, 
                           font_weight='bold',
                           font_color='black')
    
    # 添加坐标标注
    coord_labels = {node: f"({positions[node][0]:.1f}, {positions[node][1]:.1f})" 
                   for node in positions}
    
    # 调整标签位置（在节点下方）
    label_pos = {node: (positions[node][0], positions[node][1] - 0.5) 
                for node in positions}
    
    nx.draw_networkx_labels(G, label_pos, coord_labels,
                           font_size=9,
                           font_color='blue',
                           verticalalignment='top')
    
    # 设置图表属性
    plt.title(title, fontsize=16, fontweight='bold', pad=20)
    plt.axis('equal')
    plt.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
    plt.xlabel('X 坐标 (m)', fontsize=12)
    plt.ylabel('Y 坐标 (m)', fontsize=12)
    
    # 添加图例
    legend_elements = [
        plt.scatter([], [], s=200, c='lightblue', edgecolors='black', 
                   linewidths=2, label='类型 1'),
        plt.scatter([], [], s=200, c='lightgreen', edgecolors='black', 
                   linewidths=2, label='类型 2'),
        plt.scatter([], [], s=200, c='lightyellow', edgecolors='black', 
                   linewidths=2, label='类型 3'),
    ]
    plt.legend(handles=legend_elements, loc='upper right', fontsize=10)
    
    plt.tight_layout()
    plt.show()

def main():
    # 支持命令行参数
    if len(sys.argv) > 1:
        yaml_file = sys.argv[1]
    else:
        yaml_file = "floot_1_topo.yaml"
    
    print(f"📁 正在加载文件: {yaml_file}")
    
    # 加载 YAML 文件
    nodes = load_topology_yaml(yaml_file)
    
    # 解析拓扑节点
    G, positions, node_info = parse_topology_nodes(nodes)
    
    # 打印拓扑信息
    print_topology_info(G, positions, node_info)
    
    # 可视化
    title = f"拓扑地图网络 - {os.path.basename(yaml_file)}"
    visualize_topology(G, positions, node_info, title)
    
    print("\n✅ 完成！")

if __name__ == "__main__":
    main()
